#!/usr/bin/env python
# -*- coding: utf-8 -*-


import os
import json
import subprocess
import sys
from StringIO import StringIO
import requests
import PIL
import PIL.Image
from PyPDF2 import PdfFileWriter, PdfFileReader

DEFAULT_MOBILE_UA = 'Mozilla/5.0 (Linux; Android 4.3; Nexus 7 Build/JSS15Q) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/42.0.2307.2 Safari/537.36'

SWTSTRINGS = "/home/dhwei/local/bin/swfstrings"
SWTRENDER = "/home/dhwei/local/bin/swfrender"


class WenKu:
    _RN = 5
    WIDTH = 1260
    HEIGHT = 1782

    def __init__(self):
        self.url_fmt = 'http://ai.wenku.baidu.com/play/%s?pn=%d&rn=%d'
        self.outtxt = None
        self.outpdf = None
        self.width = self.WIDTH
        self.height = self.HEIGHT

    def down(self, uid, mkpdf=False):
        print 'getting %s ... ...' % (uid)

        self.get_text(uid)
        self.outtxt.close()

        if mkpdf:
            self.outpdf = PdfFileWriter()
            self.get_page(uid, 1)

            output_file = '%s.pdf' % uid
            with open(output_file, 'wb') as outputStream:
                self.outpdf.write(outputStream)

    def get_page(self, uid, pn):
        r = requests.get(self.url_fmt % (uid, pn, self._RN))
        data = r.content
        idx = data.index('   ')
        pg_info = json.JSONDecoder().decode(data[:idx])

        # print pg_info

        data = data[idx:].lstrip(' ')

        self.save_files(data, uid, pn)

        total_page = int(pg_info['totalPage'])
        next_pg = int(pg_info['toPage']) + 1

        # print next_pg, total_page

        if next_pg <= total_page:
            print 'getting %s %d/%d ... ...' % (uid, next_pg, total_page)
            self.get_page(uid, next_pg)

    def save_files(self, data, uid, pn):
        while True:
            idx = data.find('CWS', 3)
            if idx == -1:
                break

            swfname = '%s_%d.swf' % (uid, pn)
            with open(swfname, 'wb') as f:
                f.write(data[:idx])
            self.convert(uid, swfname)

            data = data[idx:]
            pn += 1

        swfname = '%s_%d.swf' % (uid, pn)
        with open(swfname, 'wb') as f:
            f.write(data)
        self.convert(uid, swfname)

    def convert(self, uid, swfname):
        # self.save_text(uid, swfname)
        pngname = self.swf_to_img(swfname)
        self.merge_to_pdf(pngname)
        os.remove(swfname)
        os.remove(pngname)

    def save_text(self, uid, swfname):
        f = open('%s.txt' % uid, 'a')
        subprocess.call([SWTSTRINGS, swfname], stdout=f)
        f.close()

    def swf_to_img(self, swfname):
        pngname = os.path.splitext(swfname)[0] + '.png'
        subprocess.call([SWTRENDER, '-X', str(self.width), '-Y', str(self.height), '-o', pngname, swfname])
        return pngname

    def merge_to_pdf(self, img):
        im = PIL.Image.open(img)

        buff = StringIO()
        im.save(buff, 'PDF', resoultion=100.0)

        input = PdfFileReader(buff)
        num_pages = input.getNumPages()

        for i in range(0, num_pages):
            self.outpdf.addPage(input.getPage(i))

    def get_text(self, uid):
        url = "http://wk.baidu.com/view/%s?pcf=2&re=view" % uid

        headers = {
            'User-Agent': DEFAULT_MOBILE_UA,
        }

        r = requests.get(url, headers=headers)
        txt = r.content
        #print txt

        flags = [ ('docinfo: ', '\n'),
                  ('docInfo: ', ',\n'),
            ]

        idx_s = -1
        idx_e = -1
        for (st, ed) in flags:
            idx_s = txt.find(st)
            if idx_s == -1:
                continue
            idx_s += len(st)
            idx_e = txt.find(ed, idx_s)
            if idx_e != -1:
                break

        if idx_s == -1 or idx_e == -1:
            print 'no docInfo'
            return False

        docinfo = txt[idx_s: idx_e]
        print docinfo
        docinfo = json.loads(docinfo)

        if not self.outtxt:
            title = docinfo.get('docInfo', {}).get('docTitle', uid)
            self.outtxt = open('%s.txt' % title, 'wb')

        md5sum = docinfo["md5sum"]
        md5sum = md5sum[1:]

        bcsParam = docinfo["bcsParam"]

        pn = 1
        range = ''
        rn = 0

        for param in bcsParam:
            rn += 1
            range += "%s_" % (param["merge"])
            if rn == 5:
                print "%d/%d" % (pn, len(bcsParam))
                range = range[:-1]
                url = 'http://retype.wenku.bdimg.com/retype/merge/%s?%s&pn=%d&rn=%d&width=176&type=org&range=%s&callback=jsonp1' % (
                uid, md5sum, pn, rn, range)
                self.get_page_text(url, uid)

                pn += rn
                rn = 0
                range = ''

        if rn:
            print "%d/%d" % (pn, len(bcsParam))
            range = range[:-1]
            url = 'http://retype.wenku.bdimg.com/retype/merge/%s?%s&pn=%d&rn=%d&width=176&type=org&range=%s&callback=jsonp1' % (
            uid, md5sum, pn, rn, range)
            self.get_page_text(url, uid)

        return True

    def get_page_text(self, url, uid):
        #print url
        headers = {
            'User-Agent': DEFAULT_MOBILE_UA,
        }

        r = requests.get(url, headers=headers)
        txt = r.content
        idx_s = len('jsonp1(')
        parags = txt[idx_s: -1]
        parags = json.loads(parags)

        for p in parags:
            # print p
            for c in p["parags"]:
                if c['t'] == 'txt':
                    self.outtxt.write(c['c'].encode('utf8'))
                    self.outtxt.write('\r\n')


if len(sys.argv) < 1:
    print "Usage: %s ID" % sys.argv[0]
    sys.exit(0)

# uid = 'a65f29777c1cfad6185fa73c'
uid = sys.argv[1]
mkpdf = False
if len(sys.argv) > 2:
    mkpdf=sys.argv[2]

wenku = WenKu()
wenku.down(uid, mkpdf)
